nanopyx.data.download

  1import os
  2import shutil
  3import tempfile
  4import zipfile
  5from urllib.request import ProxyHandler, build_opener, install_opener
  6
  7import numpy as np
  8import yaml
  9from gdown import download as gdrive_download
 10from onedrivedownloader import download as onedrive_download
 11from ..core.io.downloader import download
 12
 13from ..core.io.zip_image_loader import ZipTiffIterator
 14from .examples import get_path as get_examples_path
 15
 16
 17class ExampleDataManager:
 18    _base_path = get_examples_path()
 19    _temp_dir = os.path.join(tempfile.gettempdir(), "nanopyx_data")
 20    _to_download_path = None
 21
 22    def __init__(self, to_download_path: str = None):
 23        """
 24        Helper class for downloading example test data
 25
 26        :param to_download_path: path to download the data to. If to_download_path is None, a temporary directory
 27        will be created. Note that it will not be automatically deleted.
 28        :type to_download_path: str, optional
 29        :raises ValueError: If to_download_path is not None and does not exist
 30
 31        To clear downloads use self._clear_download()
 32        """
 33
 34        # Set download path
 35        if to_download_path is None:
 36            self._to_download_path = self._temp_dir
 37        else:
 38            self._to_download_path = to_download_path
 39
 40        # Lets check on how many examples we have available
 41        self._datasets = []
 42        for path in os.listdir(self._base_path):
 43            full_path = os.path.join(self._base_path, path)
 44            info_file_path = os.path.join(full_path, "info.yaml")
 45            if os.path.isdir(full_path) and os.path.exists(info_file_path):
 46                info_data = None
 47                with open(os.path.join(info_file_path), "r") as f:
 48                    # Load the YAML contents
 49                    info_data = yaml.load(f, Loader=yaml.FullLoader)
 50
 51                info = {
 52                    "info_path": info_file_path,
 53                    "thumbnail_path": os.path.join(
 54                        self._base_path, path, "thumbnail.jpg"
 55                    ),
 56                    "tiff_sequence_path": None,
 57                }
 58                tiff_sequence_path = os.path.join(
 59                    self._to_download_path, path, "tiff_sequence.zip"
 60                )
 61                if os.path.exists(tiff_sequence_path):
 62                    info["tiff_sequence_path"] = tiff_sequence_path
 63
 64                for key in info_data:
 65                    info[key] = info_data[key]
 66
 67                info["shape"] = tuple([int(v) for v in info["data_shape"].split(",")])
 68
 69                info["dtype"] = np.dtype(info["data_dtype"])
 70
 71                self._datasets.append(info)
 72
 73        # Fix agent
 74        proxy = ProxyHandler({})
 75        opener = build_opener(proxy)
 76        opener.addheaders = [
 77            (
 78                "User-Agent",
 79                "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_4) AppleWebKit/603.1.30"
 80                + " (KHTML, like Gecko) Version/10.1 Safari/603.1.30",
 81            )
 82        ]
 83        install_opener(opener)
 84
 85    def list_datasets(self) -> tuple:
 86        """
 87        :return: list of dataset labels
 88        """
 89        return sorted([dataset["label"] for dataset in self._datasets])
 90
 91    def list_datasets_nickname(self) -> tuple:
 92        """
 93        :return: list of dataset nicknames
 94        """
 95        return [(dataset["nickname"], dataset["label"]) for dataset in self._datasets]
 96
 97    def get_dataset_info(self, dataset_name: str) -> dict:
 98        """
 99        :param dataset_name: can be a dataset label or nickname
100        :type dataset_name: str
101        :return: dictionary with information about the dataset
102        """
103        for dataset in self._datasets:
104            if dataset_name in (dataset["label"], dataset["nickname"]):
105                return dataset
106        raise ValueError(f"{dataset_name} not found in example datasets")
107
108    def _download(self, url, file_path, download_type=None, unzip=False):
109        if os.path.exists(
110            file_path
111        ):  # or os.path.exists(os.path.splitext(file_path)[0]):
112            # raise Warning(f"already exists, no need to download: {file_path}")
113            return
114
115        if not os.path.exists(self._temp_dir):
116            os.mkdir(self._temp_dir)
117
118        base_path = os.path.split(file_path)[0]
119        if not os.path.exists(base_path):
120            os.mkdir(base_path)
121
122        if download_type == "onedrive":
123            onedrive_download(url, file_path, unzip=unzip, clean=True)
124        elif download_type == "gdrive":
125            gdrive_download(url, file_path, fuzzy=True, quiet=False)
126        else:
127            download(url, file_path)
128
129    def _copy_auxiliary_files(self, info: dict):
130        if not os.path.exists(self._to_download_path):
131            os.mkdir(self._to_download_path)
132
133        path = os.path.join(self._to_download_path, info["label"])
134        if not os.path.exists(path):
135            os.mkdir(path)
136
137        # thumbnail_path = os.path.join(path, "thumbnail.jpg")
138        # if not os.path.exists(thumbnail_path):
139        #    shutil.copyfile(info["thumbnail_path"], thumbnail_path)
140
141        info_path = os.path.join(path, "info.yaml")
142        if not os.path.exists(info_path):
143            shutil.copyfile(info["info_path"], info_path)
144
145    def download_tiff_sequence(self, dataset_name: str) -> str:
146        """
147        Downloads the tiff sequence and returns the path to the zip file
148
149        :param dataset_name: can be a dataset label or nickname
150        :type dataset_name: str
151        :return: path to the zip file
152        """
153        info = self.get_dataset_info(dataset_name)
154        path = os.path.join(self._to_download_path, info["label"])
155
156        file_path = os.path.join(path, "tiff_sequence.zip")
157        url = info["tiff_sequence_url"]
158        download_type = info["tiff_sequence_url_type"]
159
160        self._copy_auxiliary_files(info)
161        self._download(url, file_path, download_type)
162        info["tiff_sequence_path"] = file_path
163
164        return file_path
165
166    def is_downloaded(self, dataset_name: str) -> bool:
167        """
168        :param dataset_name: can be a dataset label or nickname
169        :type dataset_name: str
170        :return: True if the dataset is downloaded
171        """
172        info = self.get_dataset_info(dataset_name)
173        return info["tiff_sequence_path"] is not None
174
175    def get_ZipTiffIterator(
176        self, dataset_name: str, as_ndarray: bool = False
177    ) -> ZipTiffIterator:
178        """
179        Downloads the tiff sequence and returns the ZipTiffIterator
180
181        :param dataset_name: can be a dataset label or nickname
182        :type dataset_name: str
183        :param as_ndarray: if True, returns a numpy array instead of a ZipTiffIterator
184        :type as_ndarray: bool
185        :return: ZipTiffIterator or numpy array
186        """
187        self._show_citation_notice(dataset_name)
188        file_path = self.download_tiff_sequence(dataset_name)
189        try:
190            zti = ZipTiffIterator(file_path)
191        except zipfile.BadZipFile:
192            self.clear_downloads()
193            # try once more
194            file_path = self.download_tiff_sequence(dataset_name)
195            zti = ZipTiffIterator(file_path)
196        if not as_ndarray:
197            return zti
198        else:
199            arr = np.asarray(zti)
200            zti.close()
201            return arr
202
203    def get_thumbnail(self, dataset_name: str) -> str:
204        """
205        Returns the path to the thumbnail
206
207        :param dataset_name: can be a dataset label or nickname
208        :type dataset_name: str
209        :return: path to the thumbnail
210        """
211        info = self.get_dataset_info(dataset_name)
212        return info["thumbnail_path"]
213
214    def clear_downloads(self):
215        """
216        Deletes all downloaded datasets
217        """
218        if os.path.exists(self._temp_dir):
219            shutil.rmtree(self._temp_dir)
220
221    def _show_citation_notice(self, dataset_name: str):
222        info = self.get_dataset_info(dataset_name)
223        if info["reference"] not in [None, ""]:
224            print(
225                f"If you find the '{dataset_name}' dataset useful, please cite: "
226                + f"{info['reference']} - {info['reference_doi']}"
227            )